from solvers.solverbase import SolverBase
from solvers.nabla.nablabase import NablaBase
from solvers.stabilizations import supg
from dolfin import *
import numpy

class CrankNicolson(SolverBase,NablaBase):
    'Crank-Nicolson method with nabla_grad stabilization.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        NablaBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        V = FunctionSpace(mesh,'CG',1)
        VV = VectorFunctionSpace(mesh,'CG',1)
        DGV = FunctionSpace(mesh,'DG',0)
        DGVV = VectorFunctionSpace(mesh,'DG',0)
        DGVVV = TensorFunctionSpace(mesh,'DG',0)

        #-----------------------------------------------------------------------------------------------------------------------------
    
        def get_matrix(v,phi,matrix,cache):
            trial = TrialFunction(DGVVV)
            psi = TestFunction(DGVVV)

            if 'aVV' not in cache.keys():
                a = inner(trial,psi)*dx
                cache['aVV'] = assemble(a)

            w = conditional(gt(sqrt(inner(nabla_grad(phi),nabla_grad(phi))),1E-16),\
                            Identity(2) - outer(nabla_grad(phi),nabla_grad(phi)/inner(nabla_grad(phi),nabla_grad(phi))),\
                            Constant(0.)*Identity(2))
                                                 
            l = inner(w,psi)*dx
            b = assemble(l)
            solve(cache['aVV'],matrix.vector(),b)
            
        def get_residumm(v,phi,phi_,dt,residuum,cache):
            trial = TrialFunction(DGV)
            psi = TestFunction(DGV)

            if 'aV' not in cache.keys():
                a =  trial*psi*dx
                cache['aV'] = assemble(a)
            
            resl = (phi-phi_+dt*inner(v,nabla_grad(phi)))*psi*dx
            resb = assemble(resl)
            solve(cache['aV'],residuum.vector(),resb)                 # get the residuum
        
        #----------------------------------------------------------------------------------------------------------------------------------
       
        phi_ = interpolate(problem.phi_,V)
        phi = TrialFunction(V)
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)
        tau = Function(DGV)
        tau.vector()[:] = supg.get_tau(v,eps=supg.diffusivity)

        residuum = Function(DGV)
        matrix = Function(DGVVV)

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0
        
        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                           ) # a way to realize swith on |v| == 0.
                        
        a = phi*varphi*dx + 0.5*dt*inner(v,nabla_grad(phi))*varphi*dx +\
            tau*pert*phi*inner(v,nabla_grad(varphi))*dx +\
            0.5*dt*tau*pert*inner(dot(nabla_grad(phi),outer(v,v)),nabla_grad(varphi))*dx\
            +0.5*dt*tau*residuum*inner(dot(nabla_grad(phi),matrix),nabla_grad(varphi))*dx 

        l = phi_*varphi*dx - 0.5*dt*inner(v,nabla_grad(phi_))*varphi*dx +\
            tau*pert*phi_*inner(v,nabla_grad(varphi))*dx - 0.5*dt*tau*pert*inner(dot(nabla_grad(phi_),outer(v,v)),nabla_grad(varphi))*dx\
            -0.5*dt*tau*residuum*inner(dot(nabla_grad(phi_),matrix),nabla_grad(varphi))*dx

        t = 0
        count = 0
        maxIter = 10
        tol = 1E-3
        cache = {}

        phi = Function(V)           # solution
        phi_m = Function(V)         # solution in iteration loop
        p_ = Function(DGVV)         # projection
        while t <= T:
            t += float(dt)
            count += 1
            iter = 0
            error = 1

            get_residumm(v,phi_,phi_,dt,residuum,cache)
            get_matrix(v,phi_,matrix,cache)
            while error > tol and iter < maxIter:
                iter += 1
                A = assemble(a) 
                L = assemble(l) 
                solve(A,phi.vector(),L)
                diff = phi.vector().array() - phi_m.vector().array()
                error = numpy.linalg.norm(diff,ord=numpy.Inf)
                phi_m.assign(phi)
                get_residumm(v,0.5*(phi_m+phi_),phi_,dt,residuum,cache)
                get_matrix(v,0.5*(phi_m+phi_),matrix,cache)

            phi_.assign(phi)
            self.update(problem,phi_,t,float(dt),count)

        problem.save_data(self.saveDir,str(self.CFL))
